import numpy as np


class Matrix:
    def __init__(self, m: int, n: int):
        self.store_ = np.zeros(shape=(m, n))

    def __eq__(self, other):
        return self.store_ == other.store_

    def __repr__(self):
        return self.store_.__repr__()

    def __str__(self):
        return self.store_.__str__()

    def __getitem__(self, item):
        return self.store_[item]

    def __setitem__(self, item, val):
        self.store_[item] = val

    def __add__(self, other):
        assert self.shape == other.shape
        tmp = mat_from_array(self.store_.__add__(other.store_))
        if tmp.is_vec:
            return vec_from_mat(tmp, keep_direction=True)
        else:
            return tmp

    def __sub__(self, other):
        assert self.shape == other.shape
        tmp = mat_from_array(self.store_.__sub__(other.store_))
        if tmp.is_vec:
            return vec_from_mat(tmp, keep_direction=True)
        else:
            return tmp

    def __mul__(self, other):
        if hasattr(other, 'store_'):
            tmp = mat_from_array(self.store_.dot(other.store_))
        else:
            tmp = mat_from_array(self.store_.dot(other))
        if tmp.is_vec:
            return vec_from_mat(tmp, keep_direction=True)
        else:
            return tmp

    def __rmul__(self, other):
        if hasattr(other, 'store_'):
            tmp = mat_from_array(np.dot(other.store_, self.store_))
        else:
            tmp = mat_from_array(np.dot(other, self.store_))
        if tmp.is_vec:
            return vec_from_mat(tmp, keep_direction=True)
        else:
            return tmp

    def __neg__(self):
        return mat_from_array(self.store_.__neg__())

    def __pow__(self, power, modulo=None):
        tmp = self.store_.__pow__(power)
        if tmp.shape[0] == 1 or tmp.shape[1] == 1:
            return vec_from_array(tmp, keep_direction=True)
        else:
            return mat_from_array(tmp)

    @property
    def shape(self):
        return self.store_.shape

    @property
    def transpose(self):
        return mat_from_array(self.store_.T)

    @property
    def is_vec(self):
        return self.shape[0] == 1 or self.shape[1] == 1

    @property
    def norm(self):
        return np.sqrt(squared_norm(self))

    @property
    def squared_norm(self):
        return squared_norm(self)

    def set_value(self, item):
        self.store_ = np.array(item)

    def power(self, n: int):
        assert self.shape[0] == self.shape[1]
        return mat_from_array(np.linalg.matrix_power(self, n))


def mat_from_array(arr):
    arr = np.array(arr)
    sp = arr.shape
    if len(sp) == 2:
        tmp = Matrix(0, 0)
        tmp.store_ = arr
        return tmp
    else:
        raise (ValueError("Should be a matrix", arr.shape))


def mat_identity(n: int):
    tmp = Matrix(0, 0)
    tmp.store_ = np.identity(n)
    return tmp


class Vector(Matrix):
    """向量"""

    def __init__(self, n: int, vertical=False):
        Matrix.__init__(self, 1, n)
        if vertical:
            self.store_ = self.store_.T

    @property
    def get_vec(self):
        if self.vertical:
            return self.store_[:, 0]
        else:
            return self.store_[0]

    @property
    def v_vec(self):
        return vec_from_array(self.store_, vertical=True)

    @property
    def h_vec(self):
        return vec_from_array(self.store_, vertical=False)

    @property
    def vertical(self):
        if self.shape[0] == 1:
            return False
        elif self.shape[1] == 1:
            return True
        else:
            raise (TypeError("Not A Vector"))

    @property
    def transpose(self):
        return vec_from_array(self.store_.T, keep_direction=True)

    def __len__(self):
        return self.shape[0] + self.shape[1] - 1

    def __getitem__(self, item):
        if type(item) == int:
            return self.store_[item, 0] if self.vertical else self.store_[0, item]
        elif len(item) == 2:
            return self.store_[item]
        else:
            raise (ValueError("item should have length 1 or 2"))

    def __setitem__(self, item, val):
        if type(item) == int:
            if self.vertical:
                self.store_[item, 0] = val
            else:
                self.store_[0, item] = val
        elif len(item) == 2:
            self.store_[item] = val
        else:
            raise (ValueError("item should have length 1 or 2"))

    def __str__(self):
        return self.get_vec.__str__()

    def __repr__(self):
        return self.get_vec.__repr__()

    def __neg__(self):
        return vec_from_array(self.store_.__neg__(), keep_direction=True)


def vec_from_array(arr, keep_direction=False, vertical=False):
    arr = np.array(arr)
    sp = arr.shape
    if len(sp) == 1:
        tmp = Vector(sp[0])
        tmp.store_[0] = arr
    elif len(sp) == 2 and (sp[0] - 1) * (sp[1] - 1) == 0:
        tmp = Vector(0)
        tmp.store_ = arr
    else:
        raise (ValueError("Should be a vector", arr.shape))
    if keep_direction:
        return tmp
    if tmp.vertical != vertical:
        tmp.store_ = tmp.store_.T
    return tmp


def vec_from_mat(mat: Matrix, keep_direction=False, vertical=False):
    tmp = Vector(0)
    tmp.store_ = mat.store_
    if keep_direction:
        return tmp
    else:
        return tmp.v_vec if vertical else tmp.h_vec


def as_vec(aom, keep_direction=False, vertical=False):
    tmp = Vector(0)
    if isinstance(aom, (Matrix, Vector)):
        tmp.store_ = aom.store_
        if keep_direction:
            return tmp
        else:
            return tmp.v_vec if vertical else tmp.h_vec
    else:
        return vec_from_array(aom, keep_direction=keep_direction, vertical=vertical)


def qr(m: Matrix):
    """QR 分解"""
    # 未来改成自己的实现
    q, r = np.linalg.qr(m.store_)
    return mat_from_array(q), mat_from_array(r)


def chelosky(m: Matrix):
    """Chelosky 分解"""
    # 未来改成自己的实现
    L = np.linalg.cholesky(m.store_)
    return mat_from_array(L)


def lu(m: Matrix):
    """LU 分解"""
    # TODO
    pass


def svd(m: Matrix, hermitian: bool = False):
    # 未来改成自己的实现
    u, s, vh = np.linalg.svd(m.store_, hermitian=hermitian)
    return mat_from_array(u), mat_from_array(s), mat_from_array(vh)


def squared_norm(m: Matrix):
    return float(np.sum(m.store_ ** 2))


def a_squared_norm(x: Vector, a: Matrix):
    """a 是正定矩阵"""
    x = vec_from_mat(x)
    return float((x.h_vec * a * x.v_vec)[0, 0])
